from math import cos, sin
import numpy as np


class Module:
    x: float
    y: float
    theta: float
    v: float

    vMax: float
    vMin: float
    aMax: float
    aMin: float
    wMax: float

    def __init__(self, vMax: float, vMin: float, aMax: float, aMin: float, wMax: float):
        self.x = 0
        self.y = 0
        self.theta = 0
        self.v = 0
        self.vMax = vMax
        self.vMin = vMin
        self.aMax = aMax
        self.aMin = aMin
        self.wMax = wMax

    def update(self, w: float, a: float, dt: float) -> None:
        a = min(a, self.aMax)
        self.x = self.x + self.v * cos(self.theta) * dt
        self.y = self.y + self.v * sin(self.theta) * dt
        self.theta = self.theta + w * dt
        self.v = self.v + a * dt

    def getState(self) -> np.ndarray:
        return np.array([self.theta, self.v, self.x, self.y]).reshape(-1, 1)

    def getMatrixA(self, v: float, theta: float, dt: float) -> np.ndarray:
        return np.array(
            [
                [1, 0, 0, 0],
                [0, 1, 0, 0],
                [-v * sin(theta) * dt, cos(theta) * dt, 1, 0],
                [v * cos(theta) * dt, sin(theta) * dt, 0, 1],
            ]
        )

    def getMatrixB(self, dt: float) -> np.ndarray:
        return np.array([[1, 0], [0, 1], [0, 0], [0, 0]]) * dt

    def getStateRange(self) -> tuple[np.ndarray, np.ndarray]:
        return (
            np.array([-np.inf, self.vMin, -np.inf, -np.inf]).reshape(-1, 1),
            np.array([np.inf, self.vMax, np.inf, np.inf]).reshape(-1, 1),
        )

    def getControlRange(self) -> tuple[np.ndarray, np.ndarray]:
        return (
            np.array([-self.wMax, self.aMin]).reshape(-1, 1),
            np.array([self.wMax, self.aMax]).reshape(-1, 1),
        )
